{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<center>\n",
    "    <p style=\"text-align:center\">\n",
    "        <img alt=\"phoenix logo\" src=\"https://storage.googleapis.com/arize-phoenix-assets/assets/phoenix-logo-light.svg\" width=\"200\"/>\n",
    "        <br>\n",
    "        <a href=\"https://docs.arize.com/phoenix/\">Docs</a>\n",
    "        |\n",
    "        <a href=\"https://github.com/Arize-ai/phoenix\">GitHub</a>\n",
    "        |\n",
    "        <a href=\"https://join.slack.com/t/arize-ai/shared_invite/zt-1px8dcmlf-fmThhDFD_V_48oU7ALan4Q\">Community</a>\n",
    "    </p>\n",
    "</center>\n",
    "<h1 align=\"center\">SQL Router Query Engine Example</h1>\n",
    "\n",
    "LlamaIndex provides high-level APIs that enable users to build powerful applications in a few lines of code. However, it can be challenging to understand what is going on under the hood and to pinpoint the cause of issues. Phoenix makes your LLM applications *observable* by visualizing the underlying structure of each call to your query engine and surfacing problematic `spans`` of execution based on latency, token count, or other evaluation metrics.\n",
    "\n",
    "In this tutorial, you will:\n",
    "- Build a query engine that uses both a SQL retriever and a VectorStoreIndex using LlamaIndex\n",
    "- Record trace data in [OpenInference tracing](https://github.com/Arize-ai/openinference) format using the global `arize_phoenix` handler\n",
    "- Observe how a more complex LlamaIndex application might perform retrieval\n",
    "\n",
    "ℹ️ This notebook requires an OpenAI API key."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install Dependencies and Import Libraries\n",
    "\n",
    "Install Phoenix, LlamaIndex, and OpenAI."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install \"arize-phoenix[evals]\" \"openai>=1\" gcsfs nest-asyncio \"llama-index>=0.10.3\" \"llama-index-core\" \"openinference-instrumentation-llama-index>=1.0.0\" \"llama-index-callbacks-arize-phoenix>=0.1.2\" \"llama-index-readers-wikipedia\" \"sqlalchemy\" wikipedia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from getpass import getpass\n",
    "\n",
    "import openai\n",
    "import pandas as pd\n",
    "import phoenix as px\n",
    "import wikipedia\n",
    "from llama_index.core import Document, Settings, set_global_handler\n",
    "from llama_index.core.indices import VectorStoreIndex\n",
    "from llama_index.core.query_engine import NLSQLTableQueryEngine, RouterQueryEngine\n",
    "from llama_index.core.selectors import LLMSingleSelector\n",
    "from llama_index.core.tools import QueryEngineTool\n",
    "from llama_index.core.utilities.sql_wrapper import SQLDatabase\n",
    "from llama_index.llms.openai import OpenAI\n",
    "from sqlalchemy import (\n",
    "    create_engine,\n",
    "    text,\n",
    ")\n",
    "\n",
    "pd.set_option(\"display.max_colwidth\", 1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Launch Phoenix\n",
    "\n",
    "You can run Phoenix in the background to collect trace data emitted by any LlamaIndex application that has been instrumented with the `OpenInferenceTraceCallbackHandler`. Phoenix supports LlamaIndex's [one-click observability](https://gpt-index.readthedocs.io/en/latest/end_to_end_tutorials/one_click_observability.html) which will automatically instrument your LlamaIndex application! You can consult our [integration guide](https://docs.arize.com/phoenix/integrations/llamaindex) for a more detailed explanation of how to instrument your LlamaIndex application.\n",
    "\n",
    "Launch Phoenix and follow the instructions in the cell output to open the Phoenix UI (the UI should be empty because we have yet to run the LlamaIndex application)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "session = px.launch_app()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Enable Phoenix tracing within LlamaIndex by setting `arize_phoenix` as the global handler. This will mount Phoenix's [OpenInferenceTraceCallback](https://docs.arize.com/phoenix/integrations/llamaindex) as the global handler. Phoenix uses OpenInference traces - an open-source standard for capturing and storing LLM application traces that enables LLM applications to seamlessly integrate with LLM observability solutions such as Phoenix."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_global_handler(\"arize_phoenix\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Configure Your OpenAI API Key\n",
    "\n",
    "Set your OpenAI API key if it is not already set as an environment variable."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not (openai_api_key := os.getenv(\"OPENAI_API_KEY\")):\n",
    "    openai_api_key = getpass(\"🔑 Enter your OpenAI API key: \")\n",
    "openai.api_key = openai_api_key\n",
    "os.environ[\"OPENAI_API_KEY\"] = openai_api_key"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Prepare reference data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "First, we'll download a dataset that contains technical details of various digital cameras and convert it into an in-memory SQL database. This dataset is provided by Kaggle and more details can be found [here](https://www.kaggle.com/datasets/crawford/1000-cameras-dataset)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_info = pd.read_parquet(\n",
    "    \"https://storage.googleapis.com/arize-phoenix-assets/datasets/structured/camera-info/cameras.parquet\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "camera_info.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "engine = create_engine(\"sqlite:///:memory:\", future=True)\n",
    "camera_info.to_sql(\"cameras\", engine, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with engine.connect() as connection:\n",
    "    result = connection.execute(text(\"SELECT * FROM cameras LIMIT 5\")).all()\n",
    "\n",
    "    for row in result:\n",
    "        print(row)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, for more general queries about digital cameras, we'll download the Wikipedia page on Digital Cameras using the `wikipedia` SDK. We will convert this document into a LlamaIndex `VectorStoreIndex`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# load the Digital Camera wikipedia page\n",
    "page = wikipedia.page(pageid=52797)\n",
    "doc = Document(id_=page.pageid, text=page.content)\n",
    "\n",
    "vector_indices = []\n",
    "vector_index = VectorStoreIndex.from_documents([doc])\n",
    "vector_indices.append(vector_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Build LlamaIndex Application\n",
    "\n",
    "Let's use a simple `RouterQueryEngine` using multiple query engine tools. We will either route to the SQL retriever or the vector index built over the \"Digital Camera\" Wikipedia page."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Settings.llm = OpenAI(temperature=0.0, model=\"gpt-4\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_database = SQLDatabase(engine, include_tables=[\"cameras\"])\n",
    "\n",
    "sql_query_engine = NLSQLTableQueryEngine(\n",
    "    sql_database=sql_database,\n",
    "    tables=[\"cameras\"],\n",
    ")\n",
    "sql_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=sql_query_engine,\n",
    "    description=(\n",
    "        \"Useful for translating a natural language query into a SQL query over\"\n",
    "        \" a table containing technical details about specific digital camera models: Model,\"\n",
    "        \" Release date, Max resolution, Low resolution, Effective pixels, Zoom wide (W),\"\n",
    "        \" Zoom tele (T), Normal focus range, Macro focus range, Storage included,\"\n",
    "        \" Weight (inc. batteries), Dimensions, Price\"\n",
    "    ),\n",
    ")\n",
    "\n",
    "vector_query_engines = [index.as_query_engine() for index in vector_indices]\n",
    "vector_tools = []\n",
    "for query_engine in vector_query_engines:\n",
    "    vector_tool = QueryEngineTool.from_defaults(\n",
    "        query_engine=query_engine,\n",
    "        description=\"Useful for answering generic questions about digital cameras.\",\n",
    "    )\n",
    "    vector_tools.append(vector_tool)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_engine = RouterQueryEngine(\n",
    "    selector=LLMSingleSelector.from_defaults(),\n",
    "    query_engine_tools=([sql_tool] + vector_tools),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Make Queries and Use Phoenix to view Spans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"What is the most expensive digital camera?\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This query asked for specific details about a camera, and routed to the SQL retriever to get context for the response. The LLM-generated SQL can be seen in a Phoenix span.\n",
    "\n",
    "![A view of the Phoenix UI showing SQL retrieval](https://storage.googleapis.com/arize-phoenix-assets/assets/docs/notebooks/tracing/llama-index-sql-retrieval-tutorial/sql-retrieval.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = query_engine.query(\"Tell me about the history of digital camera sensors.\")\n",
    "print(str(response))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "More general queries are routed to the vector index.\n",
    "\n",
    "![A view of the Phoenix UI showing vector retrieval](https://storage.googleapis.com/arize-phoenix-assets/assets/docs/notebooks/tracing/llama-index-sql-retrieval-tutorial/vectorstoreindex-retrieval.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Final Thoughts\n",
    "\n",
    "LLM Traces and the accompanying OpenInference Tracing specification is designed to be a category of telemetry data that is used to understand the execution of LLMs and the surrounding application context. This is especially useful when understanding the behavior of more complex RAG applications that might make use of multiple context retrieval strategies, such as mixing a SQL retriever with more-common vector indexes.\n",
    "\n",
    "For more details on Phoenix, LLM Tracing, and LLM Evals, checkout our [documentation](https://docs.arize.com/phoenix/)."
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
